# -*- coding: utf-8 -*-
# @Author: lidongdong
# @time  : 18-12-17 上午11:13
# @file  : dataloader.py

import random
import h5py
import numpy as np


class DataLoader(object):
    """this class is used to load training data"""
    def __init__(self, h5filename, batch_size, buffer_size=0, load_all_to_memory=False):
        h = h5py.File(h5filename)

        self.image = h["image"]
        self.caption = h["caption"]
        if load_all_to_memory:
            self.image = self.image[:]
            self.caption = self.caption[:]

        self.sample_num = self.caption.shape[0]
        self.batch_size = batch_size
        self.buffer_size = buffer_size

        idx = range(self.sample_num)
        self.sample_idx = []
        for i in idx:
            for j in range(5):
                self.sample_idx.append([i,       # real image index
                                        random.sample(range(0, i) + range(i + 1, self.sample_num), 1)[0], # wrong image index
                                        [i, j]]   # caption index i: sample index j:caption subindex
                                       )

    def get_batch(self):
        """ get batch data from h5file

        Yield:
            batch: is a list which contains real_image, wrong_image and caption
        """

        random.shuffle(self.sample_idx)
        total_sample = len(self.sample_idx)
        batch_num = total_sample / self.batch_size

        for i in range(0, self.batch_size * batch_num, self.batch_size):
            batch = [[], [], []]

            for j in range(self.batch_size):
                # real_image: [64, 64, 3]  wrong_image: [64, 64, 3]  caption: [4800, ]
                real_image_idx, wrong_image_idx, (caption_idx1, caption_idx2 )= self.sample_idx[i + j]
                real_image = self.image[real_image_idx, :]
                wrong_image = self.image[wrong_image_idx, :]
                caption = self.caption[caption_idx1, caption_idx2, :]
                batch[0].append(real_image)
                batch[1].append(wrong_image)
                batch[2].append(caption)

            yield (np.asarray(batch[0]), np.asarray(batch[1]), np.asarray(batch[2]))

        batch = [[],  [],  []]
        if i < total_sample:
            for ij in range(i + j, total_sample):
                real_image_idx, wrong_image_idx, (caption_idx1, caption_idx2 )= self.sample_idx[ij]
                real_image = self.image[real_image_idx, :]
                wrong_image = self.image[wrong_image_idx, :]
                caption = self.caption[caption_idx1, caption_idx2, :]
                batch[0].append(real_image)
                batch[1].append(wrong_image)
                batch[2].append(caption)
            print "yield last {} , less than {}".format(len(batch[0]), self.batch_size)
            yield (np.asarray(batch[0]), np.asarray(batch[1]), np.asarray(batch[2]))


if __name__ == '__main__':
    import cv2
    dataloader = DataLoader("Data/images.h5", 2)
    gen = dataloader.get_batch()
    first_batch = gen.next()
    real_image = first_batch[0][1]
    real_image = real_image * 128. + 128.
    real_image = real_image.astype(np.int32)
    cv2.imwrite("real.png", real_image)
    for item in first_batch:
        print item.shape
        print np.max(item), np.min(item)
